package merge

import (
	"bytes"
	"context"
	"errors"
	"github.com/spf13/cobra"
	"github.com/yapingcat/gomedia/go-codec"
	"github.com/yapingcat/gomedia/go-mp4"
	"github.com/yapingcat/gomedia/go-mpeg2"
	"io"
	"log"
	"os"
	"path/filepath"
	"sort"
	"strconv"
	"strings"
)

var (
	InputTsDir    string
	OutputMp4Name string

	Cmd = &cobra.Command{
		Use: "merge",
		Run: func(cmd *cobra.Command, args []string) {
			fList, err := os.ReadDir(InputTsDir)
			if err != nil {
				log.Fatalln("读取目录失败", err)
				return
			}
			var tsFileList []string
			for _, f := range fList {
				if f.Type().IsRegular() && strings.HasSuffix(strings.ToLower(f.Name()), ".ts") {
					tsFileList = append(tsFileList, filepath.Join(InputTsDir, f.Name()))
				}
			}
			if len(tsFileList) == 0 {
				log.Fatalln("目录下不存在ts文件", InputTsDir)
				return
			}
			sort.Strings(tsFileList) // 按照字典顺序排序

			err = FileListToSingleMp4(cmd.Context(), MergeReq{
				TsFileList: tsFileList,
				OutputMp4:  OutputMp4Name,
			})
			if err != nil {
				log.Fatalln("合并失败", err)
				return
			}
			log.Println("合并成功", OutputMp4Name)
		},
	}
)

type MergeReq struct {
	TsFileList []string
	OutputMp4  string
}

func FileListToSingleMp4(ctx context.Context, req MergeReq) (err error) {
	mp4file, err := os.OpenFile(req.OutputMp4, os.O_CREATE|os.O_RDWR|os.O_TRUNC, 0666)
	if err != nil {
		return err
	}
	defer mp4file.Close()

	muxer, err := mp4.CreateMp4Muxer(mp4file)
	if err != nil {
		return err
	}
	var vtid uint32 // video track id
	atid := muxer.AddAudioTrack(mp4.MP4_CODEC_AAC)

	demuxer := mpeg2.NewTSDemuxer()
	var OnFrameErr error
	var audioTimestamp uint64 = 0
	aacSampleRate := -1
	demuxer.OnFrame = func(cid mpeg2.TS_STREAM_TYPE, frame []byte, pts uint64, dts uint64) {
		if OnFrameErr != nil {
			return
		}
		if cid == mpeg2.TS_STREAM_AAC {
			audioTimestamp = pts
			codec.SplitAACFrame(frame, func(aac []byte) {
				if aacSampleRate == -1 {
					adts := codec.NewAdtsFrameHeader()
					adts.Decode(aac)
					aacSampleRate = codec.AACSampleIdxToSample(int(adts.Fix_Header.Sampling_frequency_index))
				}
				err = muxer.Write(atid, aac, audioTimestamp, audioTimestamp)
				audioTimestamp += uint64(1024 * 1000 / aacSampleRate) //每帧aac采样固定为1024。aac_sampleRate 为采样率
				if err != nil {
					OnFrameErr = err
					return
				}
			})
		} else if cid == mpeg2.TS_STREAM_H264 || cid == mpeg2.TS_STREAM_H265 {
			if vtid == 0 {
				switch cid {
				case mpeg2.TS_STREAM_H264:
					vtid = muxer.AddVideoTrack(mp4.MP4_CODEC_H264)
				case mpeg2.TS_STREAM_H265:
					vtid = muxer.AddVideoTrack(mp4.MP4_CODEC_H265)
				default:
					OnFrameErr = errors.New("unknown cid2 " + strconv.Itoa(int(cid)))
					return
				}
			}
			err = muxer.Write(vtid, frame, pts, dts)
			if err != nil {
				OnFrameErr = err
				return
			}
		} else {
			OnFrameErr = errors.New("unknown cid " + strconv.Itoa(int(cid)))
			return
		}
	}
	defer func() {
		err = muxer.WriteTrailer()
		if err != nil {
			return
		}
		err = mp4file.Sync()
	}()
	for _, tsFile := range req.TsFileList {
		select {
		case <-ctx.Done():
			return ctx.Err()
		default:
		}
		var buf []byte
		buf, err = os.ReadFile(tsFile)
		if err != nil {
			return err
		}
		err = demuxer.Input(bytes.NewReader(buf))
		if err != nil && !errors.Is(err, io.ErrUnexpectedEOF) {
			return err
		}
		if OnFrameErr != nil {
			return OnFrameErr
		}

	}

	return err
}

func init() {
	Cmd.Flags().StringVarP(&InputTsDir, "input-dir", "i", "", "存放ts文件的目录(默认为当前工作目录)")
	Cmd.Flags().StringVarP(&OutputMp4Name, "output", "o", "", "输出mp4文件名(默认为输入ts文件的目录下的all.mp4)")
	cobra.MarkFlagRequired(Cmd.Flags(), "input-dir")
	cobra.MarkFlagRequired(Cmd.Flags(), "output")
}
